# cython: language_level=3
import numpy as np
from .foundation import AlgorithmConfig
import copy
from UTILS.colorful import *
from UTILS.tensor_ops import __hash__
def _flatten_helper(T, N, _tensor):
    return _tensor.view(T * N, *_tensor.size()[2:])

'''
    轨迹
'''

class TRAJ_BASE():
    key_data_type = {}
    key_data_shape = {}
    max_mem_length = -1

    def __init__(self, traj_limit, env_id):
        self.traj_limit = traj_limit
        self.env_id = env_id
        self.readonly_lock = False
        self.key_dict = []
        self.time_pointer = 0
        self.need_reward_bootstrap = False
        self.deprecated_flag = False

    def check_type_shape(self, key, first_content=None):
        if first_content is not None:
            content_type = first_content.dtype
            content_shape = first_content.shape
            if key in trajectory.key_data_type: 
                assert trajectory.key_data_type[key] == content_type
            else: 
                trajectory.key_data_type[key] = content_type
                trajectory.key_data_shape[key] = content_shape
            return content_type, content_shape
        assert key in trajectory.key_data_type
        return trajectory.key_data_type[key], trajectory.key_data_shape[key]

    def init_track(self, key, first_content):
        content = first_content
        self.check_type_shape(key, first_content)
        assert isinstance(content, np.ndarray) or isinstance(content, float), (key, content.__class__)
        # assert self.time_pointer == 0, key  # 问题
        tensor_size = ((self.traj_limit,) + tuple(content.shape))
        set_item = np.zeros(shape=tensor_size, dtype=content.dtype)
        set_item[:] = np.nan  if np.issubdtype(content.dtype, np.floating) else 0
        setattr(self, key, set_item)
        self.key_dict.append(key)

    def init_track_none(self, key):
        content_dtype, content_shape = self.check_type_shape(key)
        tensor_size = ((self.traj_limit,) + tuple(content_shape))
        set_item = np.zeros(shape=tensor_size, dtype=content_dtype)
        set_item[:] = np.nan  if np.issubdtype(content_dtype, np.floating) else 0
        setattr(self, key, set_item)
        self.key_dict.append(key)

    def remember(self, key, content):
        assert not self.readonly_lock
        if not (key in self.key_dict) and (content is not None):
            self.init_track(key=key, first_content=content)
            getattr(self, key)[self.time_pointer] = content
        elif not (key in self.key_dict) and (content is None):
            self.init_track_none(key=key)
        elif (key in self.key_dict) and (content is not None):
            getattr(self, key)[self.time_pointer] = content
        else:
            pass

    def time_shift(self):
        assert self.time_pointer < self.traj_limit
        self.time_pointer += 1

    def cut_tail(self): 
        TJ = lambda key: getattr(self, key)
        self.readonly_lock = True
        n_frame = self.time_pointer
        # check is buffer size too big
        if n_frame > trajectory.max_mem_length: 
            trajectory.max_mem_length = n_frame
            print('max_mem_length：%d, traj_limit:%d'%(trajectory.max_mem_length, self.traj_limit))
        # clip tail
        for key in self.key_dict: setattr(self, key, TJ(key)[:n_frame])


class trajectory(TRAJ_BASE):
    def __init__(self, traj_limit, env_id):
        super().__init__(traj_limit, env_id)
        self.reference_track_name = 'value_R'

    def early_finalize(self):
        assert not self.readonly_lock   # unfinished traj
        self.need_reward_bootstrap = True

    def set_terminal_obs(self, tobs):
        self.tobs = copy.deepcopy(tobs)

    def cut_tail(self): 
        # 删去多余的预留空间
        super().cut_tail()
        TJ = lambda key: getattr(self, key)
        # 进一步地， 根据这个轨迹上的NaN，删除所有无效时间点
        reference_track = getattr(self, self.reference_track_name)
        if self.need_reward_bootstrap:
            # print('need_reward_bootstrap')
            # 找到最后一个不是nan的位置
            T = np.where(~np.isnan(reference_track.squeeze()))[0][-1]
            self.boot_strap_value = {
                'bootstrap_value_R':TJ('value_R').squeeze()[T].copy(), 
                'bootstrap_value_L':TJ('value_L').squeeze()[T].copy()
            }
            assert not hasattr(self,'tobs')
            self.set_terminal_obs(TJ('g_obs')[T].copy())
            reference_track[T] = np.nan

        # deprecated if nothing in it
        p_invalid = np.isnan(reference_track).squeeze()
        p_valid = ~p_invalid
        if p_invalid.all(): #invalid traj
            self.deprecated_flag = True
            return

        # adjust reward position
        reward = TJ('reward')
        for i in reversed(range(self.time_pointer)):
            if p_invalid[i] and i != 0: # invalid, push reward forward
                reward[i-1] += reward[i]; reward[i] = np.nan
        setattr(self, 'reward', reward)
        # clip NaN
        for key in self.key_dict: setattr(self, key, TJ(key)[p_valid])
        # all done
        return

    # new finalize
    def finalize(self, hyper_reward=None):
        if hyper_reward is not None: 
            assert self.finalize
        self.readonly_lock = True
        n_frame = self.time_pointer
        assert hyper_reward is not None
        assert not self.deprecated_flag
        self.copy_track(origin_key='reward', new_key='h_reward')
        h_rewards = getattr(self, 'h_reward')

        self.gae_finalize_return(reward_key='h_reward', value_key='value_R', new_return_name='return_R')
        self.gae_finalize_return(reward_key='reward', value_key='value_L', new_return_name='return_L')
        return 

    def copy_track(self, origin_key, new_key):
        if hasattr(self, origin_key):
            origin_handle = getattr(self, origin_key)
            setattr(self, new_key, origin_handle.copy())
            new_handle = getattr(self, new_key)
            self.key_dict.append(new_key)
            #return origin_handle, new_handle
        else:
            real_key_list = [real_key for real_key in self.__dict__ if (origin_key+'>' in real_key)]
            assert len(real_key_list)>0
            for real_key in real_key_list:
                mainkey, subkey = real_key.split('>')
                self.copy_track(real_key, (new_key+'>'+subkey))
            #return

    def gae_finalize_return(self, reward_key, value_key, new_return_name):

        gamma = AlgorithmConfig.gamma # ------- gae parameters -------
        tau = AlgorithmConfig.tau
        # ------- -------------- -------
        rewards = getattr(self, reward_key)
        value = getattr(self, value_key)
        length = rewards.shape[0]
        assert rewards.shape[0]==value.shape[0]
        gae = 0
        # initalize two more tracks
        setattr(self, new_return_name, np.zeros_like(value))
        self.key_dict.append(new_return_name)

        returns = getattr(self, new_return_name)
        boot_strap = 0 if not self.need_reward_bootstrap else self.boot_strap_value['bootstrap_'+value_key]

        for step in reversed(range(length)):
            if step==(length-1): # 最后一帧
                value_preds_delta = rewards[step] + gamma * boot_strap      - value[step]
                gae = value_preds_delta
            else:
                value_preds_delta = rewards[step] + gamma * value[step + 1] - value[step]
                gae = value_preds_delta + gamma * tau * gae
            returns[step] = gae + value[step]



def calculate_sample_entropy(samples):
    key = []
    freq = []
    n_sample = len(samples)
    for s in samples:
        if s not in key:
            key.append(s)
            freq.append(1)
        else:
            i = key.index(s)
            freq[i] += 1

    entropy = 0.0
    for j,f in enumerate(freq):
        freq[j] /= n_sample
        entropy += -freq[j] * np.log(freq[j])
    # print亮红(key)
    # print亮红(freq)
    return entropy












class TrajPoolManager(object):
    def __init__(self, n_pool):
        self.n_pool =  n_pool
        # self.traj_pool_history = []
        self.hyper_reward = []
        self.traj_pool_index = []
        self.cnt = 0

        self.clip_entropy_max = 4
        self.entropy_coef = 0

    def absorb_finalize_pool(self, pool):
        # self.traj_pool_history.append(pool)   # OOM
        pattern = []
        for traj_handle in pool:
            traj_handle.cut_tail()
        # h_reward = np.array([]).mean()
        # pattern_entropy = calculate_sample_entropy(pattern)
        # print亮绿('entropy:%.3f'%pattern_entropy)
        # h_reward = min(self.clip_entropy_max, pattern_entropy)*self.entropy_coef
        h_reward = 0
        # print亮绿('h_reward:%.3f'%h_reward)
        pool = list(filter(lambda traj: not traj.deprecated_flag, pool))
        for traj_handle in pool: 
            traj_handle.finalize(hyper_reward=h_reward)
        self.cnt += 1
        if self.cnt%2==0:
            task = ['train_R']
        else:
            task = ['train_L']
        return task, pool












'''
    轨迹池管理
'''
class BatchTrajManager():
    templete = {
        # exam that trajectory have at least following things 
        "on-policy":  ['_SKIP_', 'obs', 'actions', 'reward', '_DONE_', 'value', 'actionLogProbs'],
        "off-policy": ['_SKIP_', 'obs', 'actions', 'reward', '_DONE_'],

    }
    def __init__(self, n_env, traj_limit, templete, trainer_hook, purge_outdated):
        self.trainer_hook = trainer_hook
        self.n_env = n_env
        self.traj_limit = traj_limit
        self.train_traj_needed = AlgorithmConfig.train_traj_needed
        self.upper_training_epoch = AlgorithmConfig.upper_training_epoch
        self.live_trajs = [trajectory(self.traj_limit, env_id=i) for i in range(n_env)]
        self.live_traj_frame = [0 for _ in range(self.n_env)]
        self.traj_pool = []
        self.registered_keys = []
        self._traj_lock_buf = None
        self.pool_manager = TrajPoolManager(n_pool=self.upper_training_epoch)
        self.patience = 1e3
        self.update_cnt = 0
        self.purge_outdated = purge_outdated

    def update(self, traj_frag, index):
        assert traj_frag is not None
        for j, env_i in enumerate(index):
            traj_handle = self.live_trajs[env_i]
            for key in traj_frag:
                if traj_frag[key] is None:
                    assert False, key
                if isinstance(traj_frag[key], dict):  # 如果是二重字典，特殊处理
                    for sub_key in traj_frag[key]:
                        content = traj_frag[key][sub_key][j]
                        traj_handle.remember(key + ">" + sub_key, content)
                else:
                    content = traj_frag[key][j]
                    traj_handle.remember(key, content)
            self.live_traj_frame[env_i] += 1
            traj_handle.time_shift()
        return

    # 函数入口
    def feed_traj(self, traj_frag, require_hook=False):
        assert self._traj_lock_buf is None
        # an unlock hook must be exected before new trajectory feed in
        if require_hook: # the traj_frag is not intact, lock up traj_frag, wait for more
            assert '_DONE_' not in traj_frag
            assert 'reward' not in traj_frag
            self._traj_lock_buf = traj_frag
            return self._unlock_hook
        else:
            assert '_DONE_' in traj_frag
            assert '_SKIP_' in traj_frag
            self.__batch_update(traj_frag=traj_frag)

    def get_traj_frame(self):
        return self.live_traj_frame
        
    def train_and_clear_traj_pool(self):
        print('do update %d'%self.update_cnt)

        if self.purge_outdated:
            for env_index in range(self.n_env): 
                if self.live_traj_frame[env_index] == 0: continue # 正常结束，不用分割
                self.live_trajs[env_index].early_finalize() # 提前结束，进行切割
            
                self.traj_pool.append(self.live_trajs[env_index])
                self.live_trajs[env_index] = trajectory(self.traj_limit, env_id=env_index)
                self.live_traj_frame[env_index] = 0


        current_task_l, self.traj_pool = self.pool_manager.absorb_finalize_pool(pool=self.traj_pool)
        for current_task in current_task_l:
            ppo_update_cnt = self.trainer_hook(self.traj_pool, current_task)

        self.traj_pool = []
        self.update_cnt += 1
        # assert ppo_update_cnt == self.update_cnt
        return self.update_cnt

    def can_exec_training(self):
        if len(self.traj_pool) >= self.train_traj_needed:
            return True
        else:
            return False
            
    def _unlock_hook(self, traj_frag):
        assert self._traj_lock_buf is not None
        traj_frag.update(self._traj_lock_buf)
        self._traj_lock_buf = None
        assert '_DONE_' in traj_frag
        assert '_SKIP_' in traj_frag
        self.__batch_update(traj_frag=traj_frag)

    def ___check_integraty(self, traj_frag):
        # can not alway waste time checking this
        if self.patience < 0: return
        self.patience -= 1
        for key in traj_frag:
            if key not in self.registered_keys and (not key.startswith('_')):
                self.registered_keys.append(key)
        for key in self.registered_keys:
            assert key in traj_frag, ('this key sometimes disappears from the traj_frag:', key)

    def __batch_update(self, traj_frag):
        self.___check_integraty(traj_frag)
        done = traj_frag['_DONE_']; traj_frag.pop('_DONE_') # done flag
        skip = traj_frag['_SKIP_']; traj_frag.pop('_SKIP_') # skip/frozen flag
        tobs = traj_frag['_TOBS_']; traj_frag.pop('_TOBS_') # terminal obs
        # single bool to list bool
        if isinstance(done, bool): done = [done for i in range(self.n_env)]
        if isinstance(skip, bool): skip = [skip for i in range(self.n_env)]

        # feed
        cnt = 0
        for env_i in range(self.n_env):
            if skip[env_i]: continue
            # otherwise
            frag_index = cnt; cnt += 1
            env_index = env_i
            traj_handle = self.live_trajs[env_index]
            for key in traj_frag:
                self.make_traj_remember(traj=traj_handle, key=key, content=traj_frag[key],frag_index=frag_index)
            self.live_traj_frame[env_index] += 1
            traj_handle.time_shift()
            if done[env_i]:
                assert tobs[env_i] is not None # get the final obs
                traj_handle.set_terminal_obs(tobs[env_i])
                self.traj_pool.append(traj_handle)
                self.live_trajs[env_index] = trajectory(self.traj_limit, env_id=env_index)
                self.live_traj_frame[env_index] = 0
        # print红(self.live_traj_frame)
    def make_traj_remember(self, traj, key, content, frag_index):
        if content is None:
            traj.remember(key, None)
        elif isinstance(content, dict):
            for sub_key in content: 
                self.make_traj_remember(traj, "".join((key , ">" , sub_key)), content=content[sub_key], frag_index=frag_index)
        else:
            traj.remember(key, content[frag_index])
